"""
RAG Embeddings - 向量数据库模块

该模块封装了ChromaDB向量数据库的操作和嵌入模型的集成：
- 初始化和管理ChromaDB数据库连接
- 集成中文嵌入模型（BGE-small-zh）
- 提供文档向量化存储、查询、更新、删除等操作
- 支持余弦相似度检索

Author: xiaodeme@163.com
Date: 2025-07-09 17:32:07
"""

from chromadb import PersistentClient
from chromadb.utils import embedding_functions
from sentence_transformers import SentenceTransformer

from config import Config
from document import Document


class VectorDB:
    """向量数据库管理类，封装ChromaDB操作与嵌入模型"""

    def __init__(self, db_path=None, model_name=None, collection_name=None):
        """
        初始化向量数据库连接和嵌入模型

        Args:
            db_path (str, optional): ChromaDB数据库存储路径，默认使用配置文件设置
            model_name (str, optional): 嵌入模型名称，默认使用BGE-small-zh
            collection_name (str, optional): 集合名称，默认使用配置文件设置
        """
        # 设置数据库参数，优先使用传入参数，否则使用配置文件默认值
        self.db_path = db_path or Config.CHROMA_DB_PATH
        self.model_name = model_name or Config.EMBEDDING_MODEL
        self.collection_name = collection_name or Config.COLLECTION_NAME

        # 初始化嵌入模型
        self._load_embedding_model()
        
        # 初始化ChromaDB客户端和集合
        self._initialize_database()

    def _load_embedding_model(self):
        """加载中文嵌入模型"""
        print(f"[INFO] 正在加载嵌入模型: {self.model_name}")
        try:
            # 加载本地中文嵌入模型（需要科学上网下载，国内建议使用本地缓存）
            self.model = SentenceTransformer(self.model_name, trust_remote_code=True)
            print(f"[SUCCESS] 嵌入模型加载成功")
        except Exception as e:
            print(f"[ERROR] 嵌入模型加载失败: {e}")
            raise

        # 定义自定义嵌入函数类
        class CustomEmbeddingFunction(embedding_functions.EmbeddingFunction):
            """自定义嵌入函数，集成SentenceTransformer模型"""
            
            def __init__(self, model):
                self.model = model

            def __call__(self, texts):
                """将文本转换为向量表示"""
                return self.model.encode(texts, convert_to_numpy=True).tolist()

        # 实例化自定义嵌入函数
        self.embedding_func = CustomEmbeddingFunction(model=self.model)

    def _initialize_database(self):
        """初始化ChromaDB数据库连接和集合"""
        try:
            # 创建持久化客户端并指定数据库路径
            client = PersistentClient(path=self.db_path)
            print(f"[INFO] 已连接到ChromaDB，存储路径: {self.db_path}")

            # 创建或获取集合
            self.collection = client.get_or_create_collection(
                name=self.collection_name,
                embedding_function=self.embedding_func,
                metadata={"hnsw:space": "cosine"}  # 使用余弦相似度，值越小越相似
            )
            print(f"[INFO] 已加载集合: {self.collection_name}")
            
        except Exception as e:
            print(f"[ERROR] 数据库初始化失败: {e}")
            raise

    def add_chunks(self, chunks):
        """
        添加分块文本到向量数据库

        Args:
            chunks (list[str]): 分块后的文本列表

        Returns:
            bool: 添加操作是否成功
        """
        if not chunks:
            print("[WARNING] 输入的文本块列表为空")
            return False

        try:
            # 为每个文本块生成唯一ID
            ids = [f"chunk_{i}" for i in range(len(chunks))]
            
            # 添加文档到集合中
            self.collection.add(documents=chunks, ids=ids)
            print(f"[SUCCESS] 已添加 {len(chunks)} 个文本块到向量数据库")
            return True
            
        except Exception as e:
            print(f"[ERROR] 添加文本块失败: {e}")
            return False

    def query(self, query_text, n_results=3):
        """
        查询最相关的文本块

        Args:
            query_text (str): 查询语句
            n_results (int): 返回结果数量，默认3个

        Returns:
            dict: 查询结果，包含documents和distances等信息
        """
        if not query_text:
            print("[WARNING] 查询文本为空")
            return {"documents": [[]], "distances": [[]]}

        try:
            # 执行相似度查询
            results = self.collection.query(
                query_texts=[query_text], 
                n_results=n_results
            )
            
            print(f"\n[QUERY] 查询: '{query_text}' 的相关结果:")
            if results["documents"] and results["documents"][0] and results["distances"] and results["distances"][0]:
                for i, (doc, score) in enumerate(zip(results["documents"][0], results["distances"][0])):
                    print(f"结果 {i+1} (相似度: {1-score:.4f}): {doc[:100]}...")
            else:
                print("未找到相关结果")
                
            return results
            
        except Exception as e:
            print(f"[ERROR] 查询失败: {e}")
            return {"documents": [[]], "distances": [[]]}

    def update_chunk(self, chunk_id, new_content):
        """
        更新指定ID的文本块

        Args:
            chunk_id (str): 要更新的文本块ID
            new_content (str): 新内容

        Returns:
            bool: 更新操作是否成功
        """
        try:
            self.collection.update(ids=[chunk_id], documents=[new_content])
            print(f"[SUCCESS] 已更新文本块 ID: {chunk_id}")
            return True
            
        except Exception as e:
            print(f"[ERROR] 更新文本块失败: {e}")
            return False

    def delete_chunk(self, chunk_id):
        """
        删除指定ID的文本块

        Args:
            chunk_id (str): 要删除的文本块ID

        Returns:
            bool: 删除操作是否成功
        """
        try:
            self.collection.delete(ids=[chunk_id])
            print(f"[SUCCESS] 已删除文本块 ID: {chunk_id}")
            return True
            
        except Exception as e:
            print(f"[ERROR] 删除文本块失败: {e}")
            return False

    def get_collection_info(self):
        """
        获取集合基本信息

        Returns:
            dict: 集合信息，包含文档数量等
        """
        try:
            count = self.collection.count()
            print(f"[INFO] 集合 '{self.collection_name}' 包含 {count} 个文档")
            return {"document_count": count, "collection_name": self.collection_name}
            
        except Exception as e:
            print(f"[ERROR] 获取集合信息失败: {e}")
            return {"document_count": 0, "collection_name": self.collection_name}


def main():
    """主函数：演示向量数据库的基本使用流程"""
    
    # Step 1: 从PDF提取文本并进行分块处理
    print("=== Step 1: 文档处理 ===")
    pdf_path = Config.PDF_PATH
    document = Document(pdf_path)
    
    # 提取PDF内容
    contents = document.extract_text_from_pdf()
    if not contents:
        print("[ERROR] PDF文本提取失败，程序退出")
        return
    
    # 文本分块
    chunks = document.split_text(contents)
    print(f"共切分为 {len(chunks)} 个文本块")

    # Step 2: 初始化向量数据库并添加文档
    print("\n=== Step 2: 向量数据库操作 ===")
    vector_db = VectorDB()
    
    # 添加文本块到向量数据库
    success = vector_db.add_chunks(chunks)
    if success:
        print("文档已成功添加到向量数据库")
        
        # 显示集合信息
        vector_db.get_collection_info()
    else:
        print("[ERROR] 文档添加失败")

    # Step 3: 演示查询功能
    print("\n=== Step 3: 查询演示 ===")
    test_query = "市场准入可以包括哪些方面?"
    results = vector_db.query(test_query)


if __name__ == '__main__':
    main()






